/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package cn.ac.iie.di.ban.news.server.speed;

import cn.ac.iie.di.commons.speeder.SpeedLimitTokenBucket;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.log4j.Logger;

/**
 *
 * @author zwlin
 */
public class MultiKeySpeedLimiter implements SpeedLimiter {

	private static final Logger logger = Logger.getLogger(MultiKeySpeedLimiter.class);

	private static final SpeedLimitTokenBucket endlessTokenBucket;

	static {
		endlessTokenBucket = new SpeedLimitTokenBucket(1, 1, 1) {
			@Override
			public boolean TryGetToken() {
				return true;
			}

			@Override
			public boolean GetToken() {
				return true;
			}
		};
		endlessTokenBucket.stop();
	}

	private static final long DEFAULT_CLEANUP_INTERVAL_MILLISEC = 5000;
	private static final boolean NEED_CLEANUP = false;

	long delay;//ms
	int defaultSpeedPerSec;
	ConcurrentHashMap<String, SpeedLimitTokenBucket> limitMap;
	boolean activated;
//	ScheduledExecutorService cleanupService = Executors.newScheduledThreadPool(1);
	Timer cleanTimer;

	public MultiKeySpeedLimiter(int defaultSpeedPerSec) {
		this.defaultSpeedPerSec = defaultSpeedPerSec;
		this.delay = Math.max(1, 1000 / defaultSpeedPerSec);
		this.activated = false;
		this.limitMap = new ConcurrentHashMap<>();
		this.cleanTimer = new Timer("MultiKeySpeedLimitorCleanupTimer");
	}

	public void limit(String key, int speedPerSec) {
		limitMap.put(key, new SpeedLimitTokenBucket(speedPerSec, 1, speedPerSec));
	}

	public void limit(Map<String, Integer> KeySpeedMap) {
		for (Map.Entry<String, Integer> entry : KeySpeedMap.entrySet()) {
			limit(entry.getKey(), entry.getValue());
		}
	}

	@Override
	public boolean tryGetThrough(String key) {
		if (!activated) {
			logger.warn("MultiKeySpeedLimitor stopped.");
			return false;
		} else {
			return activated && getTokenBucket(key).TryGetToken();
		}
	}

	@Override
	public void waitGetThrough(String key) {
		if (!activated) {
			logger.warn("MultiKeySpeedLimitor stopped.");
			return;
		} else {
			while (activated && !getTokenBucket(key).GetToken());
			return;
		}
	}

	private SpeedLimitTokenBucket getTokenBucket(String key) {
		if (!limitMap.containsKey(key)) {
			if (defaultSpeedPerSec < 0) {
				return endlessTokenBucket;
			}
			SpeedLimitTokenBucket tokenBucket = new SpeedLimitTokenBucket(defaultSpeedPerSec, 1, defaultSpeedPerSec);
			tokenBucket.start();
			SpeedLimitTokenBucket rtn = limitMap.putIfAbsent(key, tokenBucket);
			if (rtn != null) {
				tokenBucket.stop();
				return rtn;
			} else {
				return tokenBucket;
			}
		} else {
			return limitMap.get(key);
		}
	}

	@Override
	public void start() {
		if (!activated) {
			synchronized (this) {
				if (!activated) {
					for (SpeedLimitTokenBucket tokenBucket : limitMap.values()) {
						tokenBucket.start();
					}
					long delta = Math.max(DEFAULT_CLEANUP_INTERVAL_MILLISEC, 10 * delay);
					if (NEED_CLEANUP) {
						cleanTimer.schedule(new TimerTask() {
							@Override
							public void run() {
								try {
									logger.info("start cleanup multiKeySpeedLimitor");
									cleanup();
									logger.info("MultiKeySpeedLimitor cleaned up.");
								} catch (Exception e) {
									logger.error("known error happened when cleanup MultiKeySpeedLimitor.", e);
								}
							}
						}, delta, delta);
					}
					this.activated = true;
				}
			}
		}
	}

	@Override
	public void stop() {
		if (activated) {
			synchronized (this) {
				if (activated) {
					this.activated = false;
					cleanTimer.cancel();
					if (NEED_CLEANUP) {
						cleanTimer.cancel();
					}
					for (SpeedLimitTokenBucket i : limitMap.values()) {
						i.stop();
					}
					limitMap.clear();
				}
			}
		}
	}

	private void cleanup() {

	}

	@Override
	public boolean isActivated() {
		return activated;
	}

}
